import cv2
import pytesseract
import numpy as np
from PIL import Image, ImageEnhance


def enhance_image(image_path):
    """增强图像对比度，有助于提高OCR识别精度"""
    img = Image.open(image_path)
    enhancer = ImageEnhance.Contrast(img)  # 对比度增强
    img_enhanced = enhancer.enhance(2.0)  # 提高对比度因子
    return img_enhanced



def clean_and_join_lines(text):
    """
    清理并拼接提取的文字中的换行符。

    :param text: 提取出的文字内容
    :return: 拼接后更易读的文字内容
    """
    lines = text.split('\n')
    cleaned_lines = []
    current_line = ""

    for line in lines:
        stripped_line = line.strip()
        if stripped_line:
            if current_line:
                current_line += stripped_line
            else:
                current_line = stripped_line
        else:
            if current_line:
                cleaned_lines.append(current_line)
                current_line = ""

    # 处理最后一行
    if current_line:
        cleaned_lines.append(current_line)

    return "\n".join(cleaned_lines)

def extract_address_area(image_path):

    # 增强图像
    enhanced_image = enhance_image(image_path)
    # 使用 Tesseract 识别文字及其位置
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)

    # 遍历识别结果，查找可能的地址
    for i in range(len(data['text'])):
        if '北京' in data['text'][i]:  # 假设地址部分包含“地址”这个词
            # 获取地址的边界框
            (x, y, w, h) = (data['left'][i], data['top'][i], data['width'][i], data['height'][i])
            # 抠出地址区域
            address_area = image[y - 10:y + h + 100, x - 10:x + w + 400]

            # # 显示抠出的区域
            # cv2.imshow("Extracted Address Area", address_area)
            # cv2.waitKey(0)  # 等待按键
            # cv2.destroyAllWindows()
            # # 保存抠出的区域（可选）
            # cv2.imwrite("extracted_address_area.jpg", address_area)
            # 保存抠出的区域到本地
            cv2.imwrite("extracted_address_area.png", address_area)

            # 识别地址区域的文字
            address_text = pytesseract.image_to_string(address_area, lang='chi_sim')
            cleaned_text = clean_and_join_lines(address_text)
            print("识别出的地址文本:", cleaned_text)
            return
    print("未找到收货人地址区域")


def main():
    # 图片路径
    image_path = 'output.png'  # 替换为你的快递单图片路径
    # 提取地址区域
    extract_address_area(image_path)


if __name__ == "__main__":
    main()